#!/usr/bin/env python3
# Author: Armit
# Create Time: 2023/01/09 

# 同一条 seq 可以出现在不同的 pdb 文件里吗，它们的几何形态一致吗？
# => 可以，且它们 cfm 几乎是一致的，有的只是单链模型作为别的多链模型的一部分
# => 唯一 seq 似乎可以唯一决定 cfm 的空间结构

from preprocess_pdb import *

'''
top 10 freq seqs:
  (2362) MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG
  (1097) GIVEQCCTSICSLYQLENYCN
  (1020) MEKVQYLTRSAIRRASTIEMPQQARQKLQNLFINFCLILICLLLICIIVMLL
  (560) MSEQNNTEMTFQIQRIYTKDISFEAPNAPHVFQKDWQPEVKLDLDTASSQLADDVYEVVLRVTVTASLGEETAFLCEVQQGGIFSIAGIEGTQMAHCLGAYCPNILFPYARECITSMVSRGTFPQLNLAPVNFDALFMNYLQQQAGEGTEEHQDA
  (552) KKKPLDGEYFTLQIRGRERFEMFRELNEALELKDAQAGKEPG-
  (440) MLILTRKVGESINIGDDITITILGVSGQQVRIGINAPKDVAVHREEIYQRIQAGLTAPD
  (350) GAMGMQQGQMAGGRAITVFSPDGRLFQVEYAREAVKKGSTALGMKFANGVLLISDKKVRSRLIEQNSIEKIQLIDDYVAAVTSGLVADARVLVDFARISAQQEKVTYGSLVNIENLVKRVADQMQQYTQYGGVRPYGVSLIFAGIDQIGPRLFDCDPAGTINEYKATAIGSGKDAVVSFLEREYKENLPEKEAVTLGIKALKSSLEEGEELKAPEIASITVGNKYRIYDQEEVKKFL
  (336) AEELEEVVMGLIINSGQARSLAYAALKQAKQGDFAAAKAMMDQSRMALNEAHLVQTKLIEGDAGEGKMKVSLVLVEAQLHLMTSMLARELITELIELHEKLKA
  (300) PIVQNLQGQMVHQAISPRTLNAWVKVVEEKAFSPEVIPMFSALSEGATPQDLNTMLNTVGGHQAAMQMLKETINEEAAEWDRLHPVHAGPIAPGQMREPRGSDIAGTTSTLQEQIGWMTHNPPIPVGEIYKRWIILGLNKIVRMYSPTSILDIRQGPKEPFRDYVDRFYKTLRAEQASQEVKNWMTETLLVQNANPDCKTILKALGPGATLEEMMTACQGV
  (283) FVNQHLCGSHLVEALYLVCGERGFFYTPKT
'''

FIND_SEQ = 'MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG'


def worker(tasks:List[Path], pdb_fps:List[Path]):
  parser = PDBParser()

  for file in tasks:
    try:
      with gzip.open(file, 'rt', encoding='utf-8') as fh:
        id = file.stem[:-len('.pdb')]
        structure = parser.get_structure(id, fh)
    except:
      print(f'<< error ignore file {file}')
      continue

    for model in structure.get_models():
      for chain in model.get_chains():
        residues = chain.get_residues()
        try:
          seq = ''.join([NAME_TO_LABEL[r.resname] for r in residues])
          if seq == FIND_SEQ: pdb_fps.append(str(file))
        except KeyError: pass


def find_pdb_by_seq():
  # multiprocess tasks & workers
  tasks: List[Path] = list(walk_for_pdb())
  n_tasks_each = math.ceil(len(tasks) / N_WORKER)
  print(f'>> found {len(tasks)} pdb files')
  print(f'      assign {n_tasks_each} tasks for each thread')

  manager = Manager()
  pdb_fps: List[str] = manager.list()
  procs = [
    Process(target=worker, 
            args=(tasks[i*n_tasks_each:(i+1)*n_tasks_each], pdb_fps), 
            daemon=True) 
      for i in range(N_WORKER)
  ]
  for proc in procs: proc.start()
  for proc in procs: proc.join()
  del procs, manager, tasks

  pdb_fps = list(set(pdb_fps))
  print(f'>> found {len(pdb_fps)} items')

  with open(Path(DATA_PATH) / 'find_pdb_by_seq.txt', 'w', encoding='utf-8') as fh:
    for fp in pdb_fps:
      fh.write(fp)
      fh.write('\n')


if __name__ == '__main__':
  find_pdb_by_seq()
